package org.finesys.common.log.trace.executor;

import com.alibaba.ttl.TtlCallable;
import com.alibaba.ttl.TtlRunnable;
import org.jetbrains.annotations.NotNull;

import java.util.Collection;
import java.util.List;
import java.util.concurrent.*;

/**
 * The TtlThreadPoolExecutor class is a custom implementation of the ThreadPoolExecutor class in Java.
 * It is designed to work with the Transmittable Thread Local (TTL) library, which allows for the propagation of thread-local variables across different threads.
 */
public class TtlThreadPoolExecutor extends ThreadPoolExecutor {

    /**
     * @param corePoolSize:    The number of threads to keep in the pool, even if they are idle.
     * @param maximumPoolSize: The maximum number of threads to allow in the pool.
     * @param keepAliveTime:   When the number of threads is greater than the core, this is the maximum time that excess idle threads will wait for new tasks before terminating.
     * @param unit:            The time unit for the keepAliveTime argument.
     * @param workQueue:       The queue to use for holding tasks before they are executed.
     */
    public TtlThreadPoolExecutor(int corePoolSize, int maximumPoolSize, long keepAliveTime, @NotNull TimeUnit unit, @NotNull BlockingQueue<Runnable> workQueue) {
        super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue);
    }

    public TtlThreadPoolExecutor(int corePoolSize, int maximumPoolSize, long keepAliveTime, @NotNull TimeUnit unit, @NotNull BlockingQueue<Runnable> workQueue, @NotNull ThreadFactory threadFactory) {
        super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue, threadFactory);
    }

    public TtlThreadPoolExecutor(int corePoolSize, int maximumPoolSize, long keepAliveTime, @NotNull TimeUnit unit, @NotNull BlockingQueue<Runnable> workQueue, @NotNull RejectedExecutionHandler handler) {
        super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue, handler);
    }

    public TtlThreadPoolExecutor(int corePoolSize, int maximumPoolSize, long keepAliveTime, @NotNull TimeUnit unit, @NotNull BlockingQueue<Runnable> workQueue, @NotNull ThreadFactory threadFactory, @NotNull RejectedExecutionHandler handler) {
        super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue, threadFactory, handler);
    }

    /**
     * Executes the given command at some time in the future. The command
     * may execute in a new thread, in a pooled thread, or in the calling
     * thread, at the discretion of the <tt>Executor</tt> implementation.
     *
     * @param command the runnable task
     * @throws RejectedExecutionException if this task cannot be
     *                                    accepted for execution
     * @throws NullPointerException       if command is null
     */
    @Override
    public void execute(@NotNull Runnable command) {
        // Wrap the original Runnable with TtlRunnable to ensure that the current thread's context is preserved.
        Runnable runnable = TtlRunnable.get(command);
        // Ensure that the wrapped Runnable is not null.
        assert runnable != null;
        // Execute the wrapped Runnable using the superclass's execute method.
        super.execute(runnable);
    }

    /**
     * Submits a Runnable task for execution and returns a Future representing that task.
     * The Future's get method will return null upon successful completion.
     *
     * @param task the task to submit
     * @return a Future representing pending completion of the task
     * @throws RejectedExecutionException if the task cannot be scheduled for execution
     * @throws NullPointerException       if the task is null
     */

    @NotNull
    @Override
    public Future<?> submit(@NotNull Runnable task) {
        // Wrap the original Runnable with TtlRunnable to ensure that the current thread's context is preserved.
        Runnable runnable = TtlRunnable.get(task);
        // Ensure that the wrapped Runnable is not null.
        assert runnable != null;
        // Execute the wrapped Runnable using the superclass's execute method.
        return super.submit(runnable);

    }

    /**
     * Submits a Runnable task for execution and returns a Future representing that task.
     * The Future's get method will return the given result upon successful completion.
     *
     * @param task   the task to submit
     * @param result the result to return
     * @return a Future representing pending completion of the task
     * @throws RejectedExecutionException if the task cannot be scheduled for execution
     * @throws NullPointerException       if the task is null
     */
    @NotNull
    @Override
    public <T> Future<T> submit(@NotNull Runnable task, T result) {
        // Wrap the original Runnable with TtlRunnable to ensure that the current thread's context is preserved.
        Runnable runnable = TtlRunnable.get(task);
        // Ensure that the wrapped Runnable is not null.
        assert runnable != null;
        // Execute the wrapped Runnable using the superclass's submit method and return the Future.
        return super.submit(runnable, result);
    }

    /**
     * Submits a Callable task for execution and returns a Future representing that task.
     * The Future's get method will return the result of the Callable's call method upon successful completion.
     *
     * @param task the task to submit
     * @return a Future representing pending completion of the task
     * @throws RejectedExecutionException if the task cannot be scheduled for execution
     * @throws NullPointerException       if the task is null
     */
    @NotNull
    @Override
    public <T> Future<T> submit(@NotNull Callable<T> task) {
        // Wrap the original Callable with TtlCallable to ensure that the current thread's context is preserved.
        Callable<T> runnable = TtlCallable.get(task);
        // Ensure that the wrapped Callable is not null.
        assert runnable != null;
        // Execute the wrapped Callable using the superclass's submit method and return the Future.
        return super.submit(runnable);
    }

    /**
     * Creates a new {@link RunnableFuture} that will, upon running, execute the
     * given {@code Callable}, and arrange that {@code get} will return the
     * result of the Callable's call method upon successful completion.
     *
     * @param callable the callable task
     * @return a new {@link RunnableFuture}
     * @throws NullPointerException if the callable is null
     */
    @Override
    protected <T> RunnableFuture<T> newTaskFor(Callable<T> callable) {
        // Wrap the original Callable with TtlCallable to ensure that the current thread's context is preserved.
        Callable<T> command = TtlCallable.get(callable);
        // Create a new RunnableFuture using the superclass's newTaskFor method and return it.
        return super.newTaskFor(command);
    }

    /**
     * Creates a new {@link RunnableFuture} that will, upon running, execute the
     * given {@code Runnable}, and arrange that {@code get} will return the
     * given result on successful completion.
     *
     * @param runnable the runnable task
     * @param value    the result to return
     * @return a new {@link RunnableFuture}
     * @throws NullPointerException if the runnable is null
     */
    @Override
    protected <T> RunnableFuture<T> newTaskFor(Runnable runnable, T value) {
        // Wrap the original Runnable with TtlRunnable to ensure that the current thread's context is preserved.
        Runnable command = TtlRunnable.get(runnable);
        // Create a new RunnableFuture using the superclass's newTaskFor method and return it.
        return super.newTaskFor(command, value);
    }

    /**
     * Executes the given collection of Callable tasks, returning the result
     * of one that has completed successfully (if any), without waiting
     * for the remaining tasks to complete.
     *
     * @param tasks the collection of Callable tasks
     * @return the result returned by one of the tasks
     * @throws InterruptedException if the current thread was interrupted while waiting
     * @throws ExecutionException   if no task successfully completes
     */
    @Override
    public <T> T invokeAny(@NotNull Collection<? extends Callable<T>> tasks) throws InterruptedException, ExecutionException {
        // Wrap the original Callable tasks with TtlCallable to ensure that the current thread's context is preserved.
        List<TtlCallable<T>> callables = TtlCallable.gets(tasks);
        // Execute the wrapped Callable tasks using the superclass's invokeAny method and return the result.
        return super.invokeAny(callables);
    }

    /**
     * Executes the given collection of Callable tasks, returning the result
     * of one that has completed successfully (if any), without waiting
     * for the remaining tasks to complete.
     *
     * @param tasks   the collection of Callable tasks
     * @param timeout the maximum time to wait
     * @param unit    the time unit of the timeout argument
     * @return the result returned by one of the tasks
     * @throws InterruptedException if the current thread was interrupted while waiting
     * @throws ExecutionException   if no task successfully completes
     */
    @Override
    public <T> T invokeAny(@NotNull Collection<? extends Callable<T>> tasks, long timeout, @NotNull TimeUnit unit) throws InterruptedException, ExecutionException, TimeoutException {
        // Wrap the original Callable tasks with TtlCallable to ensure that the current thread's context is preserved.
        List<TtlCallable<T>> callables = TtlCallable.gets(tasks);
        // Execute the wrapped Callable tasks using the superclass's invokeAll method and return the result.
        return super.invokeAny(callables, timeout, unit);
    }

    /**
     * Executes the given collection of Callable tasks, returning the result
     * of one that has completed successfully (if any), without waiting
     * for the remaining tasks to complete.
     *
     * @param tasks the collection of Callable tasks
     * @return the result returned by one of the tasks
     * @throws if the current thread was interrupted while waiting
     * @throws if no task successfully completes
     * @throws if the wait timed out
     */
    @NotNull
    @Override
    public <T> List<Future<T>> invokeAll(@NotNull Collection<? extends Callable<T>> tasks) throws InterruptedException {
        // Wrap the original Callable tasks with TtlCallable to ensure that the current thread's context is preserved.
        List<TtlCallable<T>> callables = TtlCallable.gets(tasks);
        // Execute the wrapped Callable tasks using the superclass's invokeAll method and return the result.
        return super.invokeAll(callables);
    }

    /**
     * Executes the given collection of Callable tasks, returning the result
     * of one that has completed successfully (if any), without waiting
     * for the remaining tasks to complete.
     *
     * @param tasks   the collection of Callable tasks
     * @param timeout the maximum time to wait
     * @param unit    the time unit of the timeout argument
     * @return the result returned by one of the tasks
     * @throws if the current thread was interrupted while waiting
     * @throws if no task successfully completes
     * @throws if the wait timed out
     */
    @NotNull
    @Override
    public <T> List<Future<T>> invokeAll(@NotNull Collection<? extends Callable<T>> tasks, long timeout, @NotNull TimeUnit unit) throws InterruptedException {
        // Wrap the original Callable tasks with TtlCallable to ensure that the current thread's context is preserved.
        List<TtlCallable<T>> callables = TtlCallable.gets(tasks);
        // Execute the wrapped Callable tasks using the superclass's invokeAll method and return the result.
        return super.invokeAll(callables, timeout, unit);
    }
}
